# https://blog.csdn.net/qq_38808667/article/details/128052617
# https://blog.csdn.net/qq_41938789/article/details/125738584
import os

from tree_sitter import Language, Parser

from code_processing.code_type import CppCodeType, JavaCodeType, PythonCodeType, LanguageType


class CodeProcessor:
    def __init__(self, so_path='build/my-languages.so'):
        self.so_path = so_path

    def build_library(self):
        Language.build_library(
            # so文件保存位置
            self.so_path,
            # vendor文件下git clone的仓库
            [
                'vendor/tree-sitter-c',
                'vendor/tree-sitter-cpp',
                'vendor/tree-sitter-c-sharp',
                'vendor/tree-sitter-java',
                'vendor/tree-sitter-javascript',
                'vendor/tree-sitter-python',

            ]
        )

    def split_cpp_code_chunks(self, code_file_path):
        # 注意C++对应cpp，C#对应c_sharp（！这里短横线变成了下划线）
        code_parser = self.new_parser('cpp')
        cpp_code = self.read_file(code_file_path)
        tree = code_parser.parse(bytes(cpp_code, "utf8"))
        root_node = tree.root_node
        functions = []
        # 为了确定起始行
        source_code_chunks = cpp_code.split("\n")
        for child_node in root_node.children:
            if child_node.type == CppCodeType.FD.value:
                # 直接取text
                text = child_node.text
                text = text.decode('utf-8')
                # 根据行数
                function_start_line = child_node.start_point[0]
                function_end_line = child_node.end_point[0]
                # 不在同一行
                if function_start_line != function_end_line:
                    function_code = source_code_chunks[function_start_line:function_end_line + 1]
                    function_code = "\n".join(function_code)
                else:
                    function_code = source_code_chunks[function_start_line]
                if text == function_code:
                    functions.append([child_node.start_point, child_node.end_point, function_code, code_file_path])
                else:
                    raise Exception("There is a problem with cpp code chunking,please check!")
        return functions

    def split_java_code_chunks(self, code_file_path):
        code_parser = self.new_parser('java')
        source_code = self.read_file(code_file_path)
        tree = code_parser.parse(bytes(source_code, "utf8"))
        root_node = tree.root_node
        methods = []
        try:
            for child_node in root_node.children:
                if child_node.type == JavaCodeType.CD.value:
                    for method_node in child_node.children[-1].children:
                        if method_node.type == JavaCodeType.MD.value:
                            text = method_node.text
                            text = text.decode('utf-8')
                            methods.append([method_node.start_point, method_node.end_point, text, code_file_path])

                        elif method_node.type == JavaCodeType.CD.value:
                            for node in method_node.children[-1].children:
                                if node.type == JavaCodeType.MD.value:
                                    text = node.text
                                    text = text.decode('utf-8')
                                    methods.append([node.start_point, node.end_point, text, code_file_path])

        except Exception:
            raise Exception("There is a problem with java code chunking,please check!")
        return methods

    def split_python_code_chunks(self, code_file_path):
        code_parser = self.new_parser('python')
        source_code = self.read_file(code_file_path)
        tree = code_parser.parse(bytes(source_code, "utf8"))
        root_node = tree.root_node
        methods = []
        current_nodes = []
        try:
            for child_node in root_node.children:
                if child_node.type == PythonCodeType.FD.value:
                    current_nodes.append(child_node)
                elif child_node.type == PythonCodeType.CD.value:
                    for method_node in child_node.children[-1].children:
                        if method_node.type == PythonCodeType.FD.value or method_node.type == PythonCodeType.DF.value:
                            current_nodes.append(method_node)
                elif child_node.type == PythonCodeType.DF.value:
                    child_node = child_node.children[-1]
                    for method_node in child_node.children[-1].children:
                        if method_node.type == PythonCodeType.FD.value or method_node.type == PythonCodeType.DF.value:
                            current_nodes.append(method_node)
                if current_nodes:
                    for current_node in current_nodes:
                        text = current_node.text
                        text = text.decode('utf-8')
                        methods.append([current_node.start_point, current_node.end_point, text, code_file_path])
                    current_nodes = []

        except Exception:
            raise Exception("There is a problem with python code chunking,please check!")
        return methods

    def new_parser(self, language_type):
        language = Language(self.so_path, language_type)
        code_parser = Parser()
        code_parser.set_language(language)
        return code_parser

    @staticmethod
    def read_file(code_file_path):
        with open(code_file_path, 'r', encoding='utf-8') as file:
            source_code = file.read()
        return source_code

    def process_file(self, file):
        file_type = os.path.splitext(file)[-1][1:]
        if file_type == LanguageType.Cpp.value:
            code_chunks = self.split_cpp_code_chunks(file)
        elif file_type == LanguageType.Java.value:
            code_chunks = self.split_java_code_chunks(file)
        elif file_type == LanguageType.python.value:
            code_chunks = self.split_python_code_chunks(file)
        else:
            # raise Exception('暂未支持该语言{file_type}'.format(file_type=file_type))
            code_chunks = [[]]
        return code_chunks

    def process_list(self, file_list):
        code_chunk_dict = {}
        for file in file_list:
            code_chunks = self.process_file(file)
            code_chunk_dict[file] = code_chunks
        return code_chunk_dict

    def process_tree(self, file_tree):
        d2 = {}
        for key, value in file_tree.items():
            if isinstance(value, str):
                d2[key] = self.process_file(value)
            elif isinstance(value, list):
                # 如果值是列表，递归处理列表中的字典
                new_list = []
                for item in value:
                    new_list.append(self.process_tree(item))
                d2[key] = new_list
            elif isinstance(value, dict):
                # 如果值是字典，递归处理
                d2[key] = self.process_tree(value)
            else:
                raise Exception('file_tree struct error')
        return d2
